import collections
from typing import List

from sortedcontainers import SortedList


class Solution:
    def constrainedSubsetSum(self, nums: List[int], k: int) -> int:
        # 均为负数的情况
        max_val = max(nums)
        if max_val <= 0:
            return max_val

        lst = [0]
        window1 = collections.deque()
        window2 = SortedList()
        for i2, num2 in enumerate(nums):
            if num2 >= 0:
                if len(window1) != 0:
                    if len(window1) == k:
                        lst.append(window2[-1][0])
                    window1.clear()
                    window2.clear()
                    lst.append(num2)
                else:
                    lst[-1] += num2
            else:
                if len(window1) < k:
                    window1.append((num2, i2))
                    window2.add((num2, i2))
                else:
                    v = window2[-1][0] + num2
                    num1, i1 = window1.popleft()
                    window2.remove((num1, i1))
                    window1.append((v, i2))
                    window2.add((v, i2))
        else:
            if len(window1) != 0:
                lst.append(window2[-1][0])

        ans = 0
        total = 0
        for num in lst:
            total = max(0, total + num)
            ans = max(ans, total)
        return ans


if __name__ == "__main__":
    print(Solution().constrainedSubsetSum(nums=[10, 2, -10, 5, 20], k=2))  # 37
    print(Solution().constrainedSubsetSum(nums=[-1, -2, -3], k=1))  # -1
    print(Solution().constrainedSubsetSum(nums=[10, -2, -10, -5, 20], k=2))  # 23
